import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from LSTM_one.LSTM.dataset_create import create_dataset
from LSTM_one.LSTM.lstm import LSTM

def load_and_predict(province):
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 图中字体改为黑体以兼容中文
    plt.rcParams['axes.unicode_minus'] = False  # 负号显示的问题
    # 检查省份信息是否为 None，如果是则设置默认值
    if province is None:
        province = 'default_province'

    # 构建数据路径
    data_path = 'data/2206_2404新能源汽车总体销量数据_' + province + '.csv'

    # 打印数据路径
    print(data_path)
    data = pd.read_csv(data_path, encoding='latin1')

    # 归一化处理
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaled_data = scaler.fit_transform(data['sale_count'].values.reshape(-1, 1))

    # 定义模型参数
    input_size = 1
    hidden_size = 120
    num_layers = 3
    output_size = 1
    learning_rate = 0.001
    num_epochs = 800

    # 定义 LSTM_one 模型
    lstm = LSTM(input_size, hidden_size, num_layers, output_size)

    look_back = 6
    # 加载保存的模型参数
    # 检查省份信息是否为 None，如果是则设置默认值
    if province is None:
        province = 'default_province'
    model_path = 'model/lstm_model_' + province + '.pth'
    lstm.load_state_dict(torch.load(model_path))

    # 将最后的12期数据作为初始输入
    test_input = scaled_data[-look_back:]
    predictions = []

    # 预测未来六期销量，并更新下一期预测数据的输入
    for i in range(6):
        # 将数据转换为张量并添加两个新维度（batch 和 feature）
        test_input = torch.from_numpy(test_input).type(torch.Tensor)
        test_input = test_input.unsqueeze(0)  # 在第一个维度和最后一个维度增加维度1
        # 模型预测
        lstm.eval()  # 设置模型为评估模式
        with torch.no_grad():
            prediction = lstm(test_input)  # 注意：使用新的维度
        prediction = prediction.numpy()[0][0]  # 转换成NumPy数组并取第一个值
        predictions.append(prediction)
        # 更新下一期预测数据的输入
        test_input = np.concatenate([test_input.squeeze().detach().numpy()[1:], [prediction]]).reshape(-1,
                                                                                                       1)  # 去掉新增维度，并在末尾添加预测值
    # 反归一化，得到最终的预测结果
    predictions = np.array(predictions).reshape(-1, 1)
    predictions = scaler.inverse_transform(predictions)
    predictions = predictions.astype(int)
    # # 绘制预测结果图
    plt.ion()
    plt.plot(pd.date_range(start='2024-05', end='2024-11', freq='M'), predictions, label='未来值')
    plt.xlabel('时间')
    plt.ylabel('销量')
    plt.title('LSTM模型实际值与预测值比较')
    plt.legend()
    plt.show()
    plt.close(1)

    print('未来六期的销量预测值：\n', predictions)
    # 检查省份信息是否为 None，如果是则设置默认值
    if province is None:
        province = 'default_province'
    predict_path = 'predict/' + province + '省(市)新能源汽车销量预测.txt'
    with open(predict_path,'w') as file:
        for i in range(6):
            file.write('2024-0{}: '.format(i+5) + '{}'.format(predictions[i][0]) + '\n')
    print(province,'新能源汽车销量预测完毕')